# coding=UTF-8
import re
self_test = True
real_test = False

# corpus.txt
MIN_SENTENCE_NUM = 2345
MAX_SENTENCE_NUM = 2985


def check_is_num(word):
    return set(word).issubset(set(list(map(str, range(10))) + ['.']))

def check_is_date(word):
    return set(word).issubset(set(list(map(str, range(10))) + ['/']))

def extract_id(line):
    temp = line.split("logical form id=")
    temp2 = temp[1].split(">")
    return temp2[0]


def same_replace(input):
    return input.replace("\n", "").replace(",", " ").replace("'", " ").replace("’", " ")


def format_relation(raw):
    temp = raw.replace(".", " ")
    temp2 = re.sub(r"_{1,}", " ", temp)
    return temp2


def extract_question(line):
    temp = same_replace(line.split(">	")[1])
    temp2 = re.sub(r"_{1,}", " ", temp)
    temp3 = re.sub(r"\s{1,}", " ", temp2)
    return temp3

def replace_number_and_date_after_entity_removed(new_sentence):
    sentence_words = new_sentence.split(" ")
    final_sentence = ""
    for i, w in enumerate(sentence_words):
        if check_is_num(w):
            sentence_words[i] = "NUMBER"
        elif check_is_date(w):
            sentence_words[i] = "DATE"
        final_sentence += sentence_words[i] + " "
    return final_sentence.strip()

# training time
def replace_entity_in_logic_form(question,answer):
    entity_list = []
    for index, part in enumerate(answer.split('mso:')):
        if index > 0:
            entity_candidates = part.split(' ')
            for entity in filter(lambda x: not x.startswith('?'), [entity_candidates[1], entity_candidates[2]]):
                temp = same_replace(entity)
                temp2 = re.sub(r"_{1,}", " ", temp)
                temp3 = re.sub(r"\s{1,}", " ", temp2)
                entity_list.append(temp3.strip())
    for entity in entity_list:
         question = question.replace(entity, 'ENTITY')
    return question

# training time
def find_entity1and2(question,answer):
    entity_list = []
    entity2position = []
    entity2mask = {}
    if_front_is_e1 = ""

    for index, part in enumerate(answer.split('mso:')):
        if index > 0:
            raw_entity_candidates = part.split(' ')

            entity1 = raw_entity_candidates[1]
            temp = same_replace(entity1)
            temp2 = re.sub(r"_{1,}", " ", temp)
            temp3 = re.sub(r"\s{1,}", " ", temp2)

            if temp3 in question:
                entity_list.append(temp3.strip())

            entity2 = raw_entity_candidates[2]
            temp = same_replace(entity2)
            temp2 = re.sub(r"_{1,}", " ", temp)
            temp3 = re.sub(r"\s{1,}", " ", temp2)

            if temp3 in question:
                entity_list.append(temp3.strip())

    for e in entity_list:
        entity2position.append(question.find(e))
    if len(entity_list) > 1:
        if entity2position[0]>entity2position[1]:
            if_front_is_e1 = False
        else:
            if_front_is_e1 = True

    if len(entity_list)<=1:
        if len(entity_list)==0:
            return -1,-1
        entity2mask[entity_list[0]] = "ENTITY1"
        return entity2mask,-1
    else:
        entity2mask[entity_list[0]] = "ENTITY1"
        entity2mask[entity_list[1]] = "ENTITY2"
        return entity2mask,if_front_is_e1

# test time
def predict_question_entities(question):
    entity_list, tokens = [], question.split()
    index, length = 0, len(tokens)
    word_set = {'get', 'got', 'had', 'have', 'won', 'is'}
    def check_token(token):
        return token[0].isupper() or token[0].isdigit() or token.startswith(('(', '{', '.', '-', ',')) or \
               (len(token) > 1 and token[1].isupper()) or (not token.isalpha() and '-' not in token) or \
                any(ord(char) > 240 for char in token)
    def next_short_token_check(index):
        token = tokens[index]
        if token in word_set or (token in ['from', 'in', 'since'] and index + 1 < length and tokens[index+1][0].isdigit()):
            return False
        return (len(token) <= 4 and (index + 1 < length) and check_token(tokens[index+1])) or \
               (token in ['and', 'to', 'in', '&', 'of', 'at', 'with', 'for', 'on', 'from', 'de', 'pour'] and \
                (index + 1 < length) and tokens[index + 1] in ['the', 'a', 'to', 'la', 'las'] and (index + 2 < length) \
                and check_token(tokens[index + 2]))
    while index < length:
        entity = []
        while index < length and (check_token(tokens[index]) or next_short_token_check(index)):
            if len(entity) == 0 and check_token(tokens[index]) == False:
                index += 1
                continue
            entity.append(tokens[index])
            index += 1
        if not entity:
            index += 1
        else:
            entity_list.append(' '.join(entity))
    for i,e in enumerate(entity_list):
        temp = same_replace(e)
        temp2 = re.sub(r"_{1,}", " ", temp)
        temp3 = re.sub(r"\s{2,}", " ", temp2)
        entity_list[i]=temp3
    return entity_list

def update_question(question, entity_list):
    for index, entity in enumerate(entity_list, 1):
        replace_string = 'ENTITY'
        if check_is_num(entity):
            replace_string = 'NUMBER'
        elif check_is_date(entity):
            replace_string = 'DATE'
        question = question.replace(entity, replace_string)
    return question